import numpy as np
import math

blue_num = 2 # 蓝方飞机数目
red_num = 4 # 红方飞机数目
blue_env_size = 2 # 蓝方环境中的数据，这里是x，y坐标两个数据
red_env_size =2 # 红方环境中的数据，这里是x，y坐标两个数据

s_dim = red_env_size*red_num+blue_env_size*blue_num  # 状态个数，用于RL训练
a_dim = red_num  # 动作个数（一个飞机一个动作，表示飞行方向）

a_bound = math.pi  # 动作的最高值，乘以这个范围，由于方向范围为-pi到pi，神经网络输出-1到1，那么就乘以pi
delay = 10
delay_in_action = 1

state ={"units":[
{
                "ID": 2901,               # 平台编号
                "X": 0.0,          # 平台位置x轴坐标(浮点型, 单位: 米, 下同)
                "Y": 0.0,         # 平台位置y轴坐标
            },
{
                "ID": 2902,               # 平台编号
                "X": 0.0,          # 平台位置x轴坐标(浮点型, 单位: 米, 下同)
                "Y": 0.0,         # 平台位置y轴坐标
},
{
                "ID": 2903,               # 平台编号
                "X": 0.0,          # 平台位置x轴坐标(浮点型, 单位: 米, 下同)
                "Y": 0.0,         # 平台位置y轴坐标
            },
{
                "ID": 2904,               # 平台编号
                "X": 0.0,          # 平台位置x轴坐标(浮点型, 单位: 米, 下同)
                "Y": 0.0,         # 平台位置y轴坐标
            }
]
}
state_enemy ={"units":[
{
                "ID": 2901,               # 平台编号
                "X": 0.0,          # 平台位置x轴坐标(浮点型, 单位: 米, 下同)
                "Y": 0.0,         # 平台位置y轴坐标
            },
{
                "ID": 2902,               # 平台编号
                "X": 0.0,          # 平台位置x轴坐标(浮点型, 单位: 米, 下同)
                "Y": 10.0,         # 平台位置y轴坐标
            }
]
}
enemy_drones_xy = []
for index, drone in enumerate(state_enemy["units"]):
    enemy_drones_xy.append([drone['X'],drone['Y']])
print("enemy_drones_xy:", enemy_drones_xy)
def get_theta_list( unit1, unit2):
    '''
    unit1 [10.0, 0.0]
    unit2 (0.0, 0.0)
    因此转换unit2
    '''
    unit2 = list(unit2)
    d_x = unit1[0] - unit2[0]
    d_y = unit1[1] - unit2[1]
    return math.atan2(d_y, d_x)

# 计算两点之间的距离
def calculate_distance( point1, point2):
    '''
    函数说明:
    计算两个坐标点之间的距离。

    参数说明：
    point1、point2：
        目标点:[x,y]

    返回值：
        两个坐标之间的欧氏距离
    '''
    return math.sqrt((point1[0] - point2[0]) ** 2 + (point1[1] - point2[1]) ** 2)

def surrounding_points(enemy_drones_xy, num_circle, surrounding_radius):
    '''
    包围的目标是围在目标飞机的四周，可以将目标位置设置为四周的各个点位，让飞机前往距离最近的点位，只要让所有飞机飞行距离最短就可以了。
    输入：敌方飞机的坐标list：[[],[],...]  每个飞机包围目标个数：int
    输出：敌方飞机周围包围点list:[[],[],...]
    '''
    num_enemy_drones = len(enemy_drones_xy)
    angle_step = 2 * math.pi / num_circle
    points_arrive = []
    for i in range(num_enemy_drones):
        for j in range(num_circle):
            angle = j * angle_step
            x = enemy_drones_xy[i][0] + surrounding_radius * math.cos(angle)
            y = enemy_drones_xy[i][1] + surrounding_radius * math.sin(angle)
            # 判断坐标是否为0
            if abs(x) < 0.000000001:
                x = 0.0
            if abs(y) < 0.000000001:
                y = 0.0
            points_arrive.append([x, y])
    return points_arrive


def find_nearest_drone( target_point, self_drones):
    '''
    函数说明:
    计算飞机集群到目标点最短距离的飞机下标。

    参数说明：
    target_point：
        目标点:[x,y]
    self_drones ：
        可用的飞机集合，最初来自于 obs[red][units]，其中数据需要使用['X']等方法得到，并不是一个现成的列表list

    返回值：
    drones_index[0]:
        己方飞机集合中最近的飞机下标 int
    drones_xy[0]
        己方飞机集合中最近的飞机坐标 list[x,y]
    drones_id[0]
        己方飞机集合中最近的飞机ID int
    '''
    drones_with_indices = [(index, (drone['X'], drone['Y']), drone['ID']) for index, drone in enumerate(self_drones)]
    if not drones_with_indices:
        return []  # 返回一个空列表或适当的默认值，以防止空列表引发错误
    sorted_drones = sorted(drones_with_indices,
                           key=lambda item: calculate_distance(target_point, item[1]))  # 从小到大对目标点距离进行排序
    drones_index = [item[0] for item in sorted_drones]
    drones_xy = [item[1] for item in sorted_drones]
    drones_id = [item[2] for item in sorted_drones]

    return drones_index[0], drones_xy[0], drones_id[0]


def get_surrounding_points( enemy_drones_xy, num_circle, surrounding_radius):
    '''
    函数说明:
    包围的目标是围在目标飞机的四周，可以将目标位置设置为四周的各个点位，让飞机前往距离最近的点位，只要让所有飞机飞行距离最短就可以了。

    参数说明：
    enemy_drones_xy
        敌方飞机的坐标 list：[[x,y],[x,y],...]
    num_circle
        每个敌方飞机周围包围的飞机数目 int
    surrounding_radius
        包围半径 float

    返回值：
    points_arrive
        敌方飞机周围包围点list:[[x,y],[x,y],...]

    eg:
        surrounding_points([[0,0],[1,-2]],1, 10)
        [[10.0, 0.0], [11.0, -2.0]]
    '''
    num_enemy_drones = len(enemy_drones_xy)
    angle_step = 2 * math.pi / num_circle
    points_arrive = []
    for i in range(num_enemy_drones):
        for j in range(num_circle):
            angle = j * angle_step
            x = enemy_drones_xy[i][0] + surrounding_radius * math.cos(angle)
            y = enemy_drones_xy[i][1] + surrounding_radius * math.sin(angle)
            # 判断坐标是否为0
            if abs(x) < 0.000000001:
                x = 0.0
            if abs(y) < 0.000000001:
                y = 0.0
            points_arrive.append([x, y])
    return points_arrive


def calculate_surrounding_path( points_arrive, self_drones, circle_type):
    '''
    函数说明:
    从get_surrounding_points得到需要到达的目标后，此函数用于分配每一己方飞机需要前往的地点。考虑不同的包围方式，提供接口。
    前提是，己方飞机大于等于敌方飞机*包围需要的飞机
    参数说明：
    points_arrive：
        敌方飞机周围包围点list:[[],[],...]
    self_drones ：
        可用的飞机集合，最初来自于 obs[red][units]，其中数据需要使用['X']等方法得到，并不是一个现成的列表list
    circle_type ：
        str类型数据“SHORT_PATH_SELF”、“SHORT_PATH_ALL”
        "SHORT_PATH_SELF"表示己方飞单个飞机前进路程最短，如果已经有飞机前往则进入第二近。实现简单，但是容易导致并不是最近的飞机前往，并非最优解。
        "SHORT_PATH_ALL"表示己方所有飞机前进的路程最短。实现复杂，每个飞机并非前往最近点，但理论上是较好的解。

    返回值：
    Action
        己方飞机动作列表list：[[f],[f],[f]] f表示浮点数，按照正常的角坐标系
    trace_index
        己方飞机的序列
    '''
    Action = np.zeros(red_num)
    trace_index = np.zeros(red_num)
    trace_point = []
    if circle_type == "SHORT_PATH_SELF":
        units = self_drones
        for index_ene, enemy_point_circle in enumerate(points_arrive):
            if not units:
                return Action, trace_index, trace_point
            self_index, self_point, self_ID = find_nearest_drone(enemy_point_circle, units)
            for index_sel, drone in enumerate(self_drones):
                if drone['ID'] == self_ID and Action[index_sel] == 0:
                    Action[index_sel] = get_theta_list(enemy_point_circle, self_point)
                    trace_index[index_sel] = index_ene
                    trace_point.append(enemy_point_circle)
                    units = [units[i] for i in range(len(units)) if i != self_index]  # 逐渐减小可以使用的己方飞机数目

    return Action, trace_index, trace_point


points_arrive = surrounding_points(enemy_drones_xy, 2, 10)
print("points_arrive", points_arrive)
Action, trace_id, trace_point = calculate_surrounding_path(points_arrive, state["units"], "SHORT_PATH_SELF")
print("Action", Action)
print("trace_id", trace_id)
print("trace_point", trace_point)